package util

import (
	"fmt"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"strings"
)

const (
	//query table structure
	columnQuery = "SELECT COLUMN_NAME,COLUMN_COMMENT,DATA_TYPE,IS_NULLABLE,COLUMN_KEY,COLUMN_TYPE,COLUMN_DEFAULT,EXTRA " +
		"FROM information_schema.COLUMNS " +
		"WHERE TABLE_SCHEMA = ? AND TABLE_NAME =? " +
		"ORDER BY ORDINAL_POSITION"

	//query table index
	indexQuery = "SELECT TABLE_NAME,COLUMN_NAME,INDEX_NAME,SEQ_IN_INDEX,NON_UNIQUE " +
		"FROM information_schema.STATISTICS " +
		"WHERE TABLE_SCHEMA = ? AND TABLE_NAME =?"
)

var (
	defaultDataType             = "string"
	dataType        dataTypeMap = map[string]dataTypeMapping{
		"int":        (func(string) string { return "int32" }),
		"integer":    func(string) string { return "int32" },
		"smallint":   func(string) string { return "int32" },
		"mediumint":  func(string) string { return "int32" },
		"bigint":     func(string) string { return "int64" },
		"float":      func(string) string { return "float32" },
		"double":     func(string) string { return "float64" },
		"decimal":    func(string) string { return "float64" },
		"char":       func(string) string { return "string" },
		"varchar":    func(string) string { return "string" },
		"tinytext":   func(string) string { return "string" },
		"mediumtext": func(string) string { return "string" },
		"longtext":   func(string) string { return "string" },
		"binary":     func(string) string { return "[]byte" },
		"varbinary":  func(string) string { return "[]byte" },
		"tinyblob":   func(string) string { return "[]byte" },
		"blob":       func(string) string { return "[]byte" },
		"mediumblob": func(string) string { return "[]byte" },
		"longblob":   func(string) string { return "[]byte" },
		"text":       func(string) string { return "string" },
		"json":       func(string) string { return "string" },
		"enum":       func(string) string { return "string" },
		"time":       func(string) string { return "time.Time" },
		"date":       func(string) string { return "time.Time" },
		"datetime":   func(string) string { return "time.Time" },
		"timestamp":  func(string) string { return "time.Time" },
		"year":       func(string) string { return "int32" },
		"bit":        func(string) string { return "[]uint8" },
		"boolean":    func(string) string { return "bool" },
		"tinyint": func(detailType string) string {
			if strings.HasPrefix(detailType, "tinyint(1)") {
				return "bool"
			}
			return "int32"
		},
	}
)

type dataTypeMapping func(detailType string) (finalType string)

type dataTypeMap map[string]dataTypeMapping

func (m dataTypeMap) Get(dataType, detailType string) string {
	if convert, ok := m[dataType]; ok {
		return convert(detailType)
	}
	return defaultDataType
}

// DbColumn table column's info
type DbColumn struct {
	TableName     string     `gorm:"column:TABLE_NAME"`
	ColumnName    string     `gorm:"column:COLUMN_NAME"`
	ColumnComment string     `gorm:"column:COLUMN_COMMENT"`
	DataType      string     `gorm:"column:DATA_TYPE"`
	ColumnKey     string     `gorm:"column:COLUMN_KEY"`
	ColumnType    string     `gorm:"column:COLUMN_TYPE"`
	ColumnDefault string     `gorm:"column:COLUMN_DEFAULT"`
	Extra         string     `gorm:"column:EXTRA"`
	IsNullable    string     `gorm:"column:IS_NULLABLE"`
	Indexes       []*DbIndex `gorm:"-"`

	dataTypeMap map[string]func(detailType string) (dataType string) `gorm:"-"`
}

func (c *DbColumn) IsPrimaryKey() bool {
	return c != nil && c.ColumnKey == "PRI"
}

func (c *DbColumn) AutoIncrement() bool {
	return c != nil && c.Extra == "auto_increment"
}

func (c *DbColumn) SetDataTypeMap(m map[string]func(detailType string) (dataType string)) {
	c.dataTypeMap = m
}

func (c *DbColumn) GetDataType() (memberType string) {
	if mapping, ok := c.dataTypeMap[c.DataType]; ok {
		return mapping(c.ColumnType)
	}
	return dataType.Get(c.DataType, c.ColumnType)
}

// DbIndex table index info
type DbIndex struct {
	TableName  string `gorm:"column:TABLE_NAME"`
	ColumnName string `gorm:"column:COLUMN_NAME"`
	IndexName  string `gorm:"column:INDEX_NAME"`
	SeqInIndex int32  `gorm:"column:SEQ_IN_INDEX"`
	NonUnique  int32  `gorm:"column:NON_UNIQUE"`
}

func (c *DbIndex) IsPrimaryKey() bool {
	return c != nil && c.IndexName == "PRIMARY"
}

// not primary key but unique key
func (c *DbIndex) IsUnique() bool {
	return c != nil && !c.IsPrimaryKey() && c.NonUnique == 0
}

func GroupByColumn(indexList []*DbIndex) map[string][]*DbIndex {
	columnIndexMap := make(map[string][]*DbIndex, len(indexList))
	if len(indexList) == 0 {
		return columnIndexMap
	}
	for _, idx := range indexList {
		if idx == nil {
			continue
		}
		columnIndexMap[idx.ColumnName] = append(columnIndexMap[idx.ColumnName], idx)
	}
	return columnIndexMap
}

func ConnectDB(dsn string) (*gorm.DB, error) {
	if dsn == "" {
		return nil, fmt.Errorf("dsn cannot be empty")
	}
	return gorm.Open(mysql.Open(dsn))
}

func GetDbColumns(db *gorm.DB, schemaName, tableName string) ([]*DbColumn, error) {
	var resultColums []*DbColumn
	err := db.Raw(columnQuery, schemaName, tableName).Scan(&resultColums).Error
	if err != nil {
		return nil, err
	}
	var resultIndex []*DbIndex
	err = db.Raw(indexQuery, schemaName, tableName).Scan(&resultIndex).Error
	if err != nil {
		return nil, err
	}

	im := GroupByColumn(resultIndex)
	for _, c := range resultColums {
		c.Indexes = im[c.ColumnName]
	}
	return resultColums, nil
}
